/********************************************************************
 * AUTHORS: Qingkai Shi
 *
 * BEGIN DATE: March, 2013
 *
 * LICENSE: GNU GPL v3
 * 
 * COMPILE: "g++ -DDEBUG -o pigeon pigeon.cpp mincut.cpp -lstp"
 * -DDEBUG can be removed if you do not want to read the debug 
 * information.
 ********************************************************************/

#include <stdio.h>
#include <set>
#include <algorithm>
#include <vector>
#include <string.h>
#include "option.h"
#include "mincut.h"
#include "c_interface.h"

using namespace std;

struct ExprComp{
            bool operator()(const Expr &a,const Expr &b) {
                return strcmp(exprString(a), exprString(b));
            }
};

typedef struct PI{
	Expr expr;
	set<Expr, ExprComp> vars;
} PI;

int intersectionSize(set<Expr, ExprComp> * varsi, set<Expr, ExprComp> * varsj){
	int ret = 0;
	set<Expr, ExprComp>::iterator it = varsi->begin();
   	for(;it!=varsi->end();it++){
		set<Expr, ExprComp>::iterator loc = varsj->find(*it);
		if(loc!=varsj->end()){
			ret++;
		}
	}
	return ret;
}

void collectVars(Expr root, set<Expr, ExprComp> * vars){
	int degree = getDegree (root);
	if(degree!=0){
		// recursive
		for(int idx=0; idx<degree;idx++){
			Expr child = getChild(root, idx);
			collectVars(child, vars);
		}
	}else{
		//constants or variables
		if(getExprKind(root)!=BVCONST){
#ifdef DEBUG
			printf("%s\n", exprString(root));
#endif
			vars->insert(root);
		}
	}
}

void collectPI(Expr root, vector<PI> * list){
	if(getExprKind(root)==AND){
		int degree = getDegree (root);

		for(int idx=0; idx<degree;idx++){
			Expr child = getChild(root, idx);
			collectPI(child, list);
		}
        }else{
		PI * aPI = new PI;
		aPI->expr = root;

		collectVars(root, &(aPI->vars));
#ifdef DEBUG
		printf("#%d\n", (aPI->vars).size());
#endif

		list->push_back(*aPI);
	}
}

void printPartition(vector<int> part, vector<PI> piList, char * OPTION_IN_FILE, int index);

int main(int argc, char* argv[]){
	init(argc, argv);

#ifdef DEBUG
	printf("OPTION_IN_FILE: %s\n", OPTION_IN_FILE);
#endif

	VC vc = vc_createValidityChecker();
	Expr e = vc_parseExpr(vc, OPTION_IN_FILE);
	vc_printVarDecls(vc);

	vector<PI> piList;
        collectPI(e, &piList);
#ifdef DEBUG
	printf("%d\n", piList.size());
#endif

	//establish a two-dimensional matrix, in which we store the relationship graph.
	int** relationshipGraph = new int*[piList.size()];

	for(int i = 0; i<piList.size();i++){
		relationshipGraph[i] = new int[piList.size()];
	}

	// get the relationship value
	for(int i=0; i<piList.size(); i++){
		for(int j=i; j<piList.size(); j++){
			if(i==j){
				relationshipGraph[i][j]=0;
				continue;
			}
			
			set<Expr, ExprComp> varsi=piList[i].vars;
			set<Expr, ExprComp> varsj=piList[j].vars;

			int itersection_size = intersectionSize(&varsi, &varsj);
			relationshipGraph[i][j] = itersection_size;
			relationshipGraph[j][i] = itersection_size;
		}
	}

	if(OPTION_MATRIX_FILE[0]=='\0'){
		// output the matrix to stdout
		for(int i=0; i<piList.size();i++){
			for(int j=0; j<piList.size();j++){
				if(j==piList.size()-1){
					printf("%d\n", relationshipGraph[i][j]);
				}else{
					printf("%d,", relationshipGraph[i][j]);
				}
			}
		}
	}else{
#ifdef DEBUG
		printf("OPTION_MATRIX_FILE: %s\n", OPTION_MATRIX_FILE);
#endif
		FILE *output_matrix_file = fopen(OPTION_MATRIX_FILE, "w+");

		// output the matrix
		for(int i=0; i<piList.size();i++){
			for(int j=0; j<piList.size();j++){
				if(j==piList.size()-1){
					fprintf(output_matrix_file,"%d\n", relationshipGraph[i][j]);
				}else{
					fprintf(output_matrix_file,"%d,", relationshipGraph[i][j]);
				}
			}
		}
		fflush(output_matrix_file);
		fclose(output_matrix_file); 
	}

	Partition * partition = minCut(relationshipGraph, piList.size());
	vector<vector<int> >::iterator p_it = partition->begin();
	for(int index=0; p_it!=partition->end(); index++,p_it++){
		printPartition(*p_it, piList, OPTION_IN_FILE, index);
	}
}

void printPartition(vector<int> part, vector<PI> piList, char * OPTION_IN_FILE, int index){
	// get all variables used in the partition
	set<Expr, ExprComp> vars_p1;
	for(int i =0; i<part.size();i++){
		PI pi = piList[part[i]];
		set<Expr, ExprComp> vars_pi = pi.vars;
	
		set<Expr, ExprComp>::iterator it = vars_pi.begin();
   		for(;it!=vars_pi.end();it++){
			vars_p1.insert(*it);
		}
	}

	//output the var decls into a file.
	char part_file[50];
	part_file[0]='\0';
	strcat(part_file, OPTION_IN_FILE);
	strcat(part_file, ".p");

	char indexStr[2];indexStr[0]=index+'0';indexStr[1]='\0';
	strcat(part_file, indexStr);
	strcat(part_file, ".stp");

	FILE *p1_file = fopen(part_file, "w+");

	set<Expr, ExprComp>::iterator it = vars_p1.begin();
   	for(;it!=vars_p1.end();it++){
		Expr aVar = *it;

		switch(getType(aVar)) {
    		case BITVECTOR_TYPE:
			fprintf(p1_file, "%s : BITVECTOR(%d);\n", exprString(aVar), getVWidth(aVar));
		        break;
    		case ARRAY_TYPE:
			fprintf(p1_file, "%s : ARRAY BITVECTOR(%d) OF BITVECTOR(%d);\n", exprString(aVar), getIWidth(aVar), getVWidth(aVar));
	      		break;
    		case BOOLEAN_TYPE:
			fprintf(p1_file, "%s : BOOLEAN;\n", exprString(aVar));
			break;
		default:
			printf("Unsupported type: %s", exprString(aVar));exit(-1);
    		}
	}

	for(int i =0; i<part.size();i++){
		PI pi = piList[part[i]];
		Expr e = pi.expr;
		fprintf(p1_file, "ASSERT(%s);\n", exprString(e));
	}
	fprintf(p1_file, "QUERY(FALSE);\n");

	fflush(p1_file);
	fclose(p1_file); 
}
